AWS Python SDKからAthenaにクエリして結果をS3からローカルに保存

quiver

2018.01.28

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

Amazon Athena を使うと、Amazon S3内のデータに爆速でSQLで問い合わせられます。

便利ではあるのですが、管理コンソールから操作すると

SQL の記載
SQL の実行
完了まで待つ
ダウンロードリンクからSQLの実行結果を保存

と待ち時間が長く、まとめて実行した SQL が多くある場合は、少しばかり手間です。

このような問題を解決するため、API を使って SQL の実行から SQL の実行結果の保存までを API で行うスクリプトを紹介します。

実行イメージ

$ cat foo.sql
select count(*)
from bar

$ export AWS_DEFAULT_PROFILE=test

$ python athena.py foo.sql

$ ls -1
athena.log
athena.py
foo.sql     # クエリーする SQL
foo.sql.csv # クエリー結果

$ cat foo.sql.csv # クエリー結果を確認
"_col0"
"1234"

解説

主要なポイントだけかいつまんで説明します。

プログラム全体はブログ末尾を参照下さい

Amazon Athena にクエリーを投げる

Amazon Athena にクエリーを投げるには start_query_execution API を使います。

QueryString : SQL 文を指定します。
QueryExecutionContext : クエリー先のデータベース名を指定します
ResultConfiguration : クエリーの保存先を指定します。AMCから利用する場合、デフォルトは「aws-athena-query-results--」です。

result = athena.start_query_execution(
    QueryString = sql,
    QueryExecutionContext = {
        'Database': DATABASE_NAME
    },
    ResultConfiguration = {
        'OutputLocation': 's3://' + S3BUCKET_NAME,
    }
)

クエリーの実行ステータスを確認

クエリーの実行ステータスを確認するには get_query_execution API を使います。

QueryExecutionId : start_query_execution 実行時に発行される ID を指定します。

レスポンス例です。

{
   "QueryExecution": {
      "Query": "string",
      "QueryExecutionContext": {
         "Database": "string"
      },
      "QueryExecutionId": "string",
      "ResultConfiguration": {
         "EncryptionConfiguration": {
            "EncryptionOption": "string",
            "KmsKey": "string"
         },
         "OutputLocation": "string"
      },
      "Statistics": {
         "DataScannedInBytes": number,
         "EngineExecutionTimeInMillis": number
      },
      "Status": {
         "CompletionDateTime": number,
         "State": "string",
         "StateChangeReason": "string",
         "SubmissionDateTime": number
      }
   }
}

QueryExecution -> Status -> State が開始直後の RUNNING から SUCCEEDED に遷移するまでポーリングします。

サンプルプログラムでは retrying モジュールを使い exponential backoff を使いながらポーリングしています。

retrying モジュールの使い方は、次の記事を参照下さい。